#include <stdio.h>  
#include <stdlib.h>  
#include <cuda_runtime.h>  
#include <time.h>  
// CUDA核函数  
//核函数的具体实现
#define TILE_WIDTH 16 

// 共享内存法（改进）的核
// __global__ void matrixMulKernel(const float* a, const float* b, float* c, int size)
// {
// 	__shared__ float as[TILE_WIDTH][TILE_WIDTH];
// 	__shared__ float bs[TILE_WIDTH][TILE_WIDTH];
 
// 	// 索引
// 	int bx = blockIdx.x;
// 	int by = blockIdx.y;
// 	int tx = threadIdx.x;
// 	int ty = threadIdx.y;
 
// 	// 行列
// 	int row = by * TILE_WIDTH + ty;
// 	int col = bx * TILE_WIDTH + tx;
 
// 	// 遍历不同块
// 	float temp = 0.0;
// 	for (int k = 0; k < (size-1)/TILE_WIDTH + 1; k++)
// 	{
// 		// 在同一块中的不同线程会计算各自数据，并存在共享存储中互相使用
// 		if ((k * TILE_WIDTH + tx) < size && row < size)
// 		{
// 			as[ty][tx] = a[row * size + k * TILE_WIDTH + tx];
// 		}
// 		else
// 		{
// 			as[ty][tx] = 0.0f;
// 		}
// 		if ((k * TILE_WIDTH + ty) < size && col < size)
// 		{
// 			bs[ty][tx] = b[(k * TILE_WIDTH + ty) * size + col];
// 		}
// 		else
// 		{
// 			bs[ty][tx] = 0.0f;
// 		}
// 		__syncthreads(); // 等待块中其他线程同步
// 		if (row < size && col < size)
// 		{
// 			// 将对应数据计算
// 			for (int m = 0; m < TILE_WIDTH; m++)
// 				temp += as[ty][m] * bs[m][tx];
// 		}
// 		__syncthreads(); // 等待其他线程，如果不加这一句，部分线程中的as和bs会被后续循环的计算覆盖，进而导致结果错误
// 	}
// 	if (row < size && col < size)
// 	{
// 		c[row * size + col] = temp;
// 	}
// }

  
__global__ void matrixMulKernel(const float* a, const float* b, float* c, int size)  
{  
    __shared__ float as[TILE_WIDTH][TILE_WIDTH];  
    __shared__ float bs[TILE_WIDTH][TILE_WIDTH];  
  
    // 索引  
    int bx = blockIdx.x;  
    int by = blockIdx.y;  
    int tx = threadIdx.x;  
    int ty = threadIdx.y;  
  
    // 行列  
    int row = by * TILE_WIDTH + ty;  
    int col = bx * TILE_WIDTH + tx;  
  
    float temp = 0.0;  
  
    // 只遍历需要的块  
    for (int k = 0; k < (size / TILE_WIDTH); k++)  
    {  
        // 读取a到共享内存  
        if (row < size)  
        {  
            int aIndex = row * size + k * TILE_WIDTH * size + tx;  
            as[ty][tx] = (aIndex < size * size) ? a[aIndex] : 0.0f;  
        }  
  
        // 读取b到共享内存  
        if (col < size)  
        {  
            int bIndex = (k * TILE_WIDTH + ty) * size + col;  
            bs[ty][tx] = (bIndex < size * size) ? b[bIndex] : 0.0f;  
        }  
  
        // 等待块内所有线程完成读取  
        __syncthreads();  
  
        // 计算  
        if (row < size && col < size)  
        {  
            for (int m = 0; m < TILE_WIDTH; m++)  
                temp += as[ty][m] * bs[m][tx];  
        }  
        __syncthreads(); //等待其他线程计算完，因为Pvalue要用到下一个块的计算
    }  
  
    // 写入c，并在此处同步  
    if (row < size && col < size)  
    {  
        c[row * size + col] = temp;  
    }  
}

void matrixMulCPU(float* A,  float* B, float* C,int HeightA, int WidthA,int WidthB) {  
    for (int i = 0; i < HeightA; ++i) {  // 遍历矩阵A的所有行  
        for (int j = 0; j < WidthB; ++j) {  // 遍历矩阵B和C的所有列  
            C[i * WidthB + j] = 0.0f;  // 初始化C[i][j]为0  
            for (int k = 0; k < WidthA; ++k) {  // 遍历矩阵A的列和矩阵B的行  
                C[i * WidthB + j] += A[i * WidthA + k] * B[k * WidthB + j];  // 累加乘积  
            }  
        }  
    }  
}
void printMatrix(float* matrix, int width, int height) {  
    for (int y = 0; y < height; ++y) {  
        for (int x = 0; x < width; ++x) {  
            printf("%f ", matrix[y * width + x]);  
            break;
        }  
        printf("\n");
        break;  
    }  
    printf("\n");  

}

int main(){  
    // A:2048X1024 B:1024X2048
    const int HeightA = 2047;  
    const int WidthA = 1023;  
    const int WidthB = 2047;    

    // 分配CPU内存   
    float* h_A = (float*)malloc(HeightA * WidthA * sizeof(float));  
    float* h_B = (float*)malloc(WidthA * WidthB * sizeof(float));  
    float* h_C_cpu = (float*)malloc(HeightA * WidthB * sizeof(float));  
    float* h_C_gpu = (float*)malloc(HeightA * WidthB * sizeof(float));
  
    // 初始化矩阵A和B  
    for (int i = 0; i < HeightA * WidthA; ++i) h_A[i] = 7.0f;  
    for (int i = 0; i < WidthA * WidthB; ++i) h_B[i] = 1.0f;  
  
    // 分配GPU内存  
    float* d_A, *d_B, *d_C;  
    cudaMalloc((void**)&d_A, HeightA * WidthA * sizeof(float));  
    cudaMalloc((void**)&d_B, WidthA * WidthB * sizeof(float));  
    cudaMalloc((void**)&d_C, HeightA * WidthB * sizeof(float));  
  
    // 将数据从主机传输到设备  
    cudaMemcpy(d_A, h_A, HeightA * WidthA * sizeof(float), cudaMemcpyHostToDevice);  
    cudaMemcpy(d_B, h_B, WidthA * WidthB * sizeof(float), cudaMemcpyHostToDevice);  
  
    // CUDA事件计时  
    float gpu_time;  
    cudaEvent_t start, stop;  
    cudaEventCreate(&start);  
    cudaEventCreate(&stop);  
    cudaEventRecord(start, 0);  
  
    // 设置线程块和网格的大小  
    dim3 blockSize(TILE_WIDTH, TILE_WIDTH);  
    dim3 gridSize((WidthA + blockSize.x - 1) / blockSize.x, (HeightA + blockSize.y - 1) / blockSize.y);  
    matrixMulKernel<<<gridSize, blockSize>>>(d_A, d_B, d_C,HeightA);  
  
    // 等待GPU操作完成  
    cudaDeviceSynchronize();  
    cudaEventRecord(stop, 0);  
    cudaEventSynchronize(stop);  
    cudaEventElapsedTime(&gpu_time, start, stop);  
  
    // 将结果从设备传输回主机  
    cudaMemcpy(h_C_gpu, d_C, HeightA * WidthB * sizeof(float), cudaMemcpyDeviceToHost);  
  
    // CPU计时  
    clock_t start_cpu = clock();  
  
    // CPU矩阵相乘  
    matrixMulCPU(h_A,h_B,h_C_cpu,HeightA,WidthA,WidthB);
  
    clock_t end_cpu = clock();  
    double cpu_time = (double)(end_cpu - start_cpu) / CLOCKS_PER_SEC;  
    printMatrix(h_C_cpu,WidthB,HeightA);
    printMatrix(h_C_gpu,WidthB,HeightA);

  
    // 清理  
    cudaFree(d_A);  
    cudaFree(d_B);  
    cudaFree(d_C);  
    free(h_A);  
    free(h_B);  
    free(h_C_cpu);  
    free(h_C_gpu);  
    // 输出时间  
    printf("GPU time: %f ms\n", gpu_time);  
    printf("CPU time: %f ms\n", cpu_time*1000);  
    printf("加速比: %f\n",cpu_time*1000/gpu_time);
  
    // 销毁CUDA事件  
    cudaEventDestroy(start);  
    cudaEventDestroy(stop);  
  
    return 0;  
}